package ocr

import (
	"log"
	"path/filepath"
	"time"

	"github.com/sirupsen/logrus"
	"gocv.io/x/gocv"
)

type DBDetector struct {
	*PaddleModel
	preProcess  DetPreProcess
	postProcess DetPostProcess
}

func NewDBDetector(modelSource string, config *Config) *DBDetector {
	maxSideLen := config.DetMaxSideLen
	thresh := config.DetDbThresh           // getFloat64(args, "det_db_thresh", 0.3)
	boxThresh := config.DetDbBoxThresh     //getFloat64(args, "det_db_box_thresh", 0.5)
	unClipRatio := config.DetDbUnclipRatio //getFloat64(args, "det_db_unclip_ratio", 2.0)

	detector := &DBDetector{
		PaddleModel: NewPaddleModel("det", config),
		preProcess:  NewDBProcess(make([]int, 0), maxSideLen),
		postProcess: NewDBPostProcess(thresh, boxThresh, unClipRatio),
	}
	if checkModelExists(modelSource) {
		modelSource, _ = downloadModel(filepath.Join(config.ModleDir, "det"), modelSource)
	} else {
		log.Panicf("det model path: %v not exist! Please check!", modelSource)
	}
	detector.LoadModel(modelSource)
	return detector
}

func (det *DBDetector) Run(img gocv.Mat) [][][]int {
	oriH := img.Rows()
	oriW := img.Cols()
	data, resizeH, resizeW := det.preProcess.Run(img)

	logrus.WithField("resizeH", resizeH).WithField("resizeW", resizeW).WithField("this", det).Info("resize image")

	st := time.Now()

	det.inputs[det.inNames[0]].Reshape([]int32{1, 3, int32(resizeH), int32(resizeW)})
	det.inputs[det.inNames[0]].CopyFromCpu(data)

	det.predictor.Run()

	outData := make([]float32, numElements(det.outputs[det.outNames[0]].Shape()))
	det.outputs[det.outNames[0]].CopyToCpu(outData)

	ratioH, ratioW := float64(resizeH)/float64(oriH), float64(resizeW)/float64(oriW)
	boxes := det.postProcess.Run(det.outputs[det.outNames[0]], oriH, oriW, ratioH, ratioW)
	log.Println("det_box num: ", len(boxes), ", time elapse: ", time.Since(st))
	return boxes
}
